#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据增强示例脚本
展示如何使用谷歌翻译进行数据增强
"""

import os
import sys
import pandas as pd
import argparse

# 添加项目根目录到系统路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

# 导入数据增强工具
from utils.data_augmentation import augment_data_with_translation, combine_original_and_augmented
from utils.google_translate import translate_from_ch_to_en, translate_from_en_to_cn
from proxy_manager import init_proxy_pool

def test_translation():
    """测试翻译功能"""
    print("===== 测试翻译功能 =====")
    
    # 初始化代理池
    init_proxy_pool(max_use_count=5, min_pool_size=10)
    
    # 测试中文翻译为英文
    chinese_text = "股价连续上涨，市场情绪高涨，看好后市表现"
    print(f"原始中文: {chinese_text}")
    
    english_text = translate_from_ch_to_en(chinese_text, use_proxy=False)
    print(f"翻译为英文: {english_text}")
    
    # 测试英文翻译回中文
    back_chinese_text = translate_from_en_to_cn(english_text, use_proxy=False)
    print(f"翻译回中文: {back_chinese_text}")
    
    print("\n")

def sample_augmentation_demo():
    """演示对少量样本的数据增强"""
    print("===== 少量样本数据增强示例 =====")
    
    # 创建一个简单的测试数据集
    data = {
        'post_id': [1, 2, 3, 4, 5, 6, 7, 8],
        'post_title': [
            "股价连续上涨，市场情绪高涨，看好后市表现",
            "公司业绩不及预期，投资者信心受挫",
            "公司公告称将继续加大研发投入，暂时没有分红计划",
            "市场对公司新产品发布会反应积极，股价上涨",
            "公司管理层调整，市场观望情绪浓厚",
            "行业竞争加剧，公司面临市场份额下滑的风险",
            "公司财报好于市场预期，投资者情绪高涨",
            "经济环境不确定性增加，投资者持谨慎态度"
        ],
        'sentiment': ["正面", "负面", "中性", "正面", "中性", "负面", "正面", "中性"],
        'description': [
            "涨势明显，市场乐观", 
            "业绩不佳，信心下降", 
            "研发投入增加，无分红计划",
            "新产品受欢迎，股价上涨",
            "管理层变动，市场观望",
            "竞争加剧，市场份额风险",
            "业绩超预期，投资者乐观",
            "经济不确定，投资谨慎"
        ],
        'isvaild': [True, True, True, True, True, True, True, True]
    }
    
    # 创建DataFrame
    df = pd.DataFrame(data)
    
    # 保存为临时CSV文件
    temp_file = "temp_test_data.csv"
    df.to_csv(temp_file, index=False)
    print(f"创建了临时测试数据文件: {temp_file}")
    
    # 对测试数据进行增强
    try:
        print("开始对测试数据进行增强...")
        
        # 初始化代理池
        init_proxy_pool(max_use_count=5, min_pool_size=10)
        
        # 使用多线程进行数据增强处理
        augmented_df = augment_data_with_translation(
            input_file=temp_file,
            output_file="temp_augmented_data.csv",
            sleep_time_range=(0.3, 0.8),  # 缩短休眠时间，仅用于演示
            num_workers=4  # 使用4个线程
        )
        
        # 合并原始数据和增强数据
        combined_df = combine_original_and_augmented(
            original_file=temp_file,
            augmented_file="temp_augmented_data.csv",
            output_file="temp_combined_data.csv"
        )
        
        # 展示结果
        print("\n===== 增强数据结果 =====")
        if len(augmented_df) > 0:
            print(augmented_df[['post_id', 'post_title', 'sentiment']])
        else:
            print("没有生成增强数据")
        
        print("\n===== 合并数据结果 =====")
        if 'combined_df' in locals() and len(combined_df) > 0:
            print(combined_df[['post_id', 'post_title', 'sentiment']])
        else:
            print("没有合并数据")
        
    except Exception as e:
        print(f"数据增强过程中出错: {str(e)}")
    
    finally:
        # 清理临时文件
        for file in [temp_file, "temp_augmented_data.csv", "temp_combined_data.csv"]:
            if os.path.exists(file):
                os.remove(file)
                print(f"删除临时文件: {file}")
    
    print("\n")

def real_data_augmentation():
    """实际数据增强示例"""
    print("===== 实际数据增强示例 =====")
    
    # 检查实际数据文件是否存在
    data_file = "../data/taged_data_by_llm_with_isvaild.csv"
    if not os.path.exists(data_file):
        print(f"数据文件 {data_file} 不存在，请检查路径。")
        return
    
    print(f"将处理实际数据文件: {data_file}")
    
    # 初始化代理池
    init_proxy_pool(max_use_count=5, min_pool_size=16)
    
    # 进行数据增强，使用多线程处理
    try:
        augmented_df = augment_data_with_translation(
            input_file=data_file,
            output_file="../data/taged_data_by_llm_with_isvaild_augmented_copy1.csv",
            sample_ratio=1.0,  # 只处理30%的数据用于演示
            random_seed=42,
            sleep_time_range=(0.5, 1.0),
            num_workers=8  # 使用8个线程
        )
        
        print(f"成功生成了 {len(augmented_df)} 条增强数据。")
        
        # 展示部分结果
        if len(augmented_df) > 0:
            print("\n增强数据示例:")
            for i, row in augmented_df.head(3).iterrows():
                print(f"原始情感: {row['sentiment']}")
                print(f"增强文本: {row['post_title']}")
                print("-" * 50)
        
        # 合并原始数据和增强数据
        combined_df = combine_original_and_augmented(
            original_file=data_file,
            augmented_file="../data/taged_data_by_llm_with_isvaild_augmented_copy1.csv",
            output_file="../data/taged_data_by_llm_with_isvaild_combined_copy1.csv"
        )

    except Exception as e:
        print(f"实际数据增强过程中出错: {str(e)}")

def main():
    """主函数，根据命令行参数选择要执行的示例"""
    parser = argparse.ArgumentParser(description="数据增强示例")
    parser.add_argument('--mode', type=str, choices=['test', 'sample', 'real'], default='test',
                       help='运行模式: test(测试翻译), sample(样本增强), real(实际数据增强)')
    parser.add_argument('--num_workers', type=int, default=4,
                       help='线程数量，默认为4')
    
    args = parser.parse_args()
    
    print("数据增强工具示例\n")
    
    # 根据模式选择执行的示例
    if args.mode == 'test':
        test_translation()
    elif args.mode == 'sample':
        sample_augmentation_demo()
    elif args.mode == 'real':
        real_data_augmentation()


    
    print("\n示例结束。如需处理完整数据集，请使用命令行工具。")
    print("建议命令:")
    print("python utils/data_augmentation.py --input data/taged_data_by_llm_with_isvaild.csv --combine --sample_ratio 0.5 --num_workers 8")

if __name__ == "__main__":
    # main()
    real_data_augmentation()
    # data_file = "../data/taged_data_by_llm_with_isvaild.csv"
    # combine_original_and_augmented(
    #     original_file=data_file,
    #     augmented_file="../data/taged_data_by_llm_with_isvaild_augmented.csv",
    #     output_file="../data/taged_data_by_llm_with_isvaild_combined.csv"
    # )